package com.vr.xxaiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.vr.xxaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * 抽象基础代理类
 * 用于管理代理状态和执行的抽象基类
 * 为状态转换、内存管理和基于步骤的执行循环提供基本功能。
 * 子类必须实现‘ step ’方法
 *
 * @author hzh
 * @date 2025/06/08
 */
@Data
@Slf4j
public abstract class BaseAgent {
    /**
     * 核心属性
     */
    private String name;
    /**
     * 提示
     */
    private String systemPrompt;
    private String nextStepPrompt;

    /**
     * 执行控制
     */
    private int currentStep;
    private int maxStep;

    /**
     * 大模型
     */
    private ChatClient llm;
    /**
     * 状态,默认是休眠状态
     */
    private AgentState state=AgentState.IDLE;
    /**
     * memory(需要自主维护会话上下文)
     * */
    private List<Message> messageList=new ArrayList<>();

    /**
     * 运行代理(同步)
     *
     * @param userPrompt 用户提示词
     * @return {@code String } 执行结果
     */
    public String run(String userPrompt) {
        // 只有处于休眠状态时才能运行
        if (this.state!=AgentState.IDLE) {
            throw new RuntimeException("Cannot run agent from state:"+this.state);
        }
        // 用户提示不能为空
        if (StrUtil.isBlank(userPrompt)) {
            throw new RuntimeException("Cannot run agent with empty user prompt");
        }
        // 更改状态，进入运行状态
        this.state=AgentState.RUNNING;
        // 记录上下文
        this.messageList.add(new UserMessage(userPrompt));
        // 保存结果列表
        List<String> results=new ArrayList<>();
        try {
            while(this.currentStep<this.maxStep && this.state!=AgentState.FINISHED) {
                // 新增一步
                this.currentStep+=1;
                log.info("Executing step {}/{}",  this.currentStep, this.maxStep);
                // 单步执行,拿到结果
                String stepResult = this.step();
                results.add(String.format("Step %d: %s", this.currentStep, stepResult));
            }
            // 判断步骤是否超出限制
            if (this.currentStep>=this.maxStep) {
                this.currentStep=0;
                // 手动完成
                this.state=AgentState.FINISHED;
                // 添加终止提示词，大模型收到后会自动调用终止工具(有的前提下)
                results.add(String.format("Terminated: Reached max steps (%d)", this.maxStep));
            }
            // 每个对话结果用换行符分割
            return String.join("\n", results);
        } catch (Exception e) {
            this.state=AgentState.ERROR;
            return String.format("Error: %s", e.getMessage());
        } finally {
            // 清理资源
            this.cleanup();
        }
    }

    /**
     * 运行代理(流的方式)
     *
     * @param userPrompt 用户提示词
     * @return {@code SseEmitter }
     */
    public SseEmitter runStream(String userPrompt) {
        // 创建一个sseEmitter,建立长连接，5分钟超时
        SseEmitter sseEmitter = new SseEmitter(5 * 1000 * 60L);
        // 用异步的方式，方式主线程阻塞
        CompletableFuture.runAsync(() -> {
            try {
                // 只有处于休眠状态时才能运行
                if (this.state!=AgentState.IDLE) {
                    sseEmitter.send("Cannot run agent from state:"+this.state);
                    // 完成，释放资源
                    sseEmitter.complete();
                }
                // 用户提示不能为空
                if (StrUtil.isBlank(userPrompt)) {
                    sseEmitter.send("Cannot run agent with empty user prompt");
                    // 完成，释放资源
                    sseEmitter.complete();
                }
                // 更改状态，进入运行状态
                this.state=AgentState.RUNNING;
                // 记录上下文
                this.messageList.add(new UserMessage(userPrompt));
                try {
                    while(this.currentStep<this.maxStep && this.state!=AgentState.FINISHED) {
                        // 新增一步
                        this.currentStep+=1;
                        log.info("Executing step {}/{}",  this.currentStep, this.maxStep);
                        // 单步执行,拿到结果
                        String stepResult = this.step();
                        sseEmitter.send(String.format("Step %d: %s", this.currentStep, stepResult));
                    }
                    // 判断步骤是否超出限制
                    if (this.currentStep>=this.maxStep) {
                        this.currentStep=0;
                        // 手动完成
                        this.state=AgentState.FINISHED;
                        // 添加终止提示词，大模型收到后会自动调用终止工具(有的前提下)
                        sseEmitter.send(String.format("Terminated: Reached max steps (%d)", this.maxStep));
                    }
                    // 正常完成
                    sseEmitter.complete();
                } catch (Exception e) {
                    this.state=AgentState.ERROR;
                    sseEmitter.send(String.format("Error: %s", e.getMessage()));
                    // 完成，释放资源
                    sseEmitter.complete();
                } finally {
                    // 清理资源
                    this.cleanup();
                }
            } catch (Exception e) {
                log.error("Error: {}", e.getMessage());
                sseEmitter.completeWithError(e);
            }
        });
        // 超时时的回调
        sseEmitter.onTimeout(()->{
            this.state=AgentState.ERROR;
            this.cleanup();
            log.error("SSE connect time out");
        });
        // 完成时的回调
        sseEmitter.onCompletion(()->{
            if (this.state==AgentState.RUNNING) {
                this.state=AgentState.FINISHED;
            }
            this.cleanup();
            log.info("SSE connect completed");
        });

        return sseEmitter;
    }

    /**
     * 清理资源
     */
    protected void cleanup() {
        // 子类可以重写这个方法清理资源
    }

    /**
     * 在代理的工作流中执行一个步骤。
     * 必须由子类实现来定义特定的行为
     */
    protected abstract String step();
}
